"""
线程池
"""

from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
from functools import wraps

from my_typing import *

# 线程池配置
CPU_THREAD_COUNT: Final[int] = 4
IO_THREAD_COUNT: Final[int] = 8

cpu_executor: ThreadPoolExecutor = ThreadPoolExecutor(
    max_workers=CPU_THREAD_COUNT,
    thread_name_prefix="CPU-THREAD"
)
""" 计算密集型（数据计算） 线程池 """

io_executor: ThreadPoolExecutor = ThreadPoolExecutor(
    max_workers=IO_THREAD_COUNT,
    thread_name_prefix="IO-THREAD"
)
""" IO 密集型（数据库读写） 线程池 """


def shutdown_thread_pools(wait: bool = True) -> None:
    """
    优雅地关闭所有线程池

    Args:
        wait: 是否等待所有线程执行完成
    """

    cpu_executor.shutdown(wait=wait)
    io_executor.shutdown(wait=wait)


def async_io_task(func: CallableNone) -> CallableNone:
    """ 异步任务装饰器（IO 密集型） """

    @wraps(func)
    def wrapper(*args, **kwargs) -> None:
        from components.context_manager import stash_request_ctx
        from components.global_param import get_request_context

        # 捕获当前上下文
        current_ctx: dict = deepcopy(get_request_context())

        # 创建一个新函数，使用上下文管理器
        def task_with_context(*task_args, **task_kwargs):
            with stash_request_ctx(current_ctx):
                return func(*task_args, **task_kwargs)

        # 提交带有上下文的任务
        io_executor.submit(task_with_context, *args, **kwargs)

    return wrapper


def async_cpu_task(func: CallableNone) -> CallableNone:
    """ 异步任务装饰器（CPU 密集型） """

    @wraps(func)
    def wrapper(*args, **kwargs) -> None:
        from components.context_manager import stash_request_ctx
        from components.global_param import get_request_context

        # 捕获当前上下文
        current_ctx: dict = deepcopy(get_request_context())

        # 创建一个新函数，使用上下文管理器
        def task_with_context(*task_args, **task_kwargs):
            with stash_request_ctx(current_ctx):
                return func(*task_args, **task_kwargs)

        # 提交带有上下文的任务
        cpu_executor.submit(task_with_context, *args, **kwargs)

    return wrapper
